import numpy as np
import struct
import os # 导入os模块来处理文件路径

def read_int8_matrix(filename, rows=16, cols=16):
    """
    读取int8矩阵文件（每行一个字节的十六进制），并将其解析为带符号的numpy数组。
    """
    # 检查文件是否存在
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None

    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int8_values = []
        for hex_str in hex_lines:
            byte = int(hex_str, 16)
            # 处理负数补码 (例如 0xFF -> -1)
            if byte >= 0x80:
                byte -= 0x100
            int8_values.append(byte)
        # 根据指定尺寸重塑矩阵
        return np.array(int8_values, dtype=np.int8).reshape(rows, cols)

def read_int8_matrix_col_major(filename, rows=16, cols=16):
    """
    按列优先读取int8矩阵文件（每行一个字节的十六进制），并将其解析为带符号的numpy数组。
    """
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None

    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int8_values = []
        for hex_str in hex_lines:
            byte = int(hex_str, 16)
            if byte >= 0x80:
                byte -= 0x100
            int8_values.append(byte)
        # 列优先存储
        return np.array(int8_values, dtype=np.int8).reshape(rows, cols, order='F')

def read_int32_matrix_from_big_endian_hex(filename, rows=16, cols=16):
    """
    直接读取大端int32矩阵文件（每行一个32位十六进制），并进行有符号转换。
    """
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None

    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int32_values = []
        for hex_str in hex_lines:
            # 直接读取32位十六进制值，并进行有符号转换
            int32_val = int(hex_str, 16)
            if int32_val >= 0x80000000:
                int32_val = int32_val - 0x100000000
            int32_values.append(int32_val)
        
        # 按行优先存储
        return np.array(int32_values, dtype=np.int32).reshape(rows, cols)

def saturating_add_int32(a, b):
    """
    对两个int32的numpy数组执行饱和加法, 以精确模拟硬件行为。
    """
    # 获取int32的最小和最大值
    min_val = np.iinfo(np.int32).min
    max_val = np.iinfo(np.int32).max
    
    # 将输入数组提升到64位整数以防止Python/Numpy在计算时发生环绕溢出
    a_64 = a.astype(np.int64)
    b_64 = b.astype(np.int64)
    
    # 执行64位加法，此时中间结果是精确的
    sum_64 = a_64 + b_64
    
    # 使用np.clip函数将结果钳位在32位整数的表示范围内
    clipped_sum = np.clip(sum_64, min_val, max_val)
    
    # 将最终结果转换回int32类型
    return clipped_sum.astype(np.int32)

# --- 主程序 ---

# --- 请在这里修改您的文件路径 ---
BASE_PATH = 'D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int8_int32/m8n32k16/'
FILE_A = os.path.join(BASE_PATH, 'a_int8_m8n32k16_hex.txt')
FILE_B = os.path.join(BASE_PATH, 'b_int8_m8n32k16_hex.txt')
FILE_C = os.path.join(BASE_PATH, 'c_int32_big_endian_m8n32k16_hex.txt')
OUTPUT_FILE = os.path.join(BASE_PATH, 'd_int32_m8n32k16_hex.txt') # 输出文件名

# m16n16k16矩阵尺寸
M_SIZE = 8  # A矩阵行数
N_SIZE = 32  # B矩阵列数
K_SIZE = 16  # 内积维度

print("开始执行硬件行为模拟 (m16n16k16)...")

# 1. 读取输入矩阵
print(f"正在读取矩阵 A: {FILE_A}")
A = read_int8_matrix(FILE_A, M_SIZE, K_SIZE)  # A: 8x16

print(f"正在读取矩阵 B: {FILE_B}")
B = read_int8_matrix_col_major(FILE_B, K_SIZE, N_SIZE)  # B: 16x32 (列优先)

print(f"正在读取矩阵 C: {FILE_C}")
C = read_int32_matrix_from_big_endian_hex(FILE_C, M_SIZE, N_SIZE)  # C: 8x32

# 检查文件是否读取成功
if A is None or B is None or C is None:
    print("错误: 一个或多个输入文件无法读取，程序终止。")
else:
    # 2. 计算 A * B (矩阵乘法)
    # 将A和B提升到int32进行乘法，防止中间结果溢出
    print("正在计算矩阵乘法 (A * B)...")
    print(f"A矩阵形状: {A.shape}")
    print(f"B矩阵形状: {B.shape}")
    AB = np.matmul(A.astype(np.int32), B.astype(np.int32), dtype=np.int32)
    print(f"AB矩阵形状: {AB.shape}")

    # 3. 计算 AB + C (使用饱和加法)
    print("正在执行饱和加法 (AB + C)...")
    print(f"C矩阵形状: {C.shape}")
    result = saturating_add_int32(AB, C)
    print(f"结果矩阵形状: {result.shape}")

    # 4. 将结果保存为大端格式的十六进制文件
    print(f"正在将结果写入文件: {OUTPUT_FILE}")
    with open(OUTPUT_FILE, 'w') as f:
        for row in result:
            for val in row:
                # 使用'>'表示大端模式, 'i'表示4字节有符号整数
                packed = struct.pack('>i', val)
                # 将打包后的字节转换为十六进制字符串
                hex_str = packed.hex()
                f.write(hex_str + '\n')

    print("\n模拟完成！")
    print(f"生成的黄金参考模型 (Golden Model) 已保存至: {OUTPUT_FILE}")
    print(f"结果矩阵尺寸: {M_SIZE}x{N_SIZE} = {M_SIZE*N_SIZE} 个元素")